import base64
import json
import math
import os

from PyQt5.QtCore import Qt, QRectF
from PyQt5.QtGui import QPixmap, QMouseEvent, QKeySequence, QColor

from frames.mainFrame import Ui_MainWindow
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox, QAction, QMenu, QInputDialog, \
    QShortcut, QWidget, QDialog, QGraphicsScene, QSizePolicy, QGraphicsPixmapItem
from PyQt5 import QtGui
import cv2
import ddddocr
import hashlib
from loguru import logger
from utils.YoloBody import YOLOSystem
import numpy as np
from dataCollect.WidgetFunctionStarter import DataCollectWidgetFunctionStarter
from frames.chooseConfigDialog import Ui_ChooseConfigDialog
from frames.ResizableRect import ResizableRect
from utils.DatasetsCutUtil import DatasetsCutUtil


def CvMatToQImage(cvMat):
    if len(cvMat.shape) == 2:
        # 灰度图是单通道，所以需要用Format_Indexed8
        rows, columns = cvMat.shape
        bytesPerLine = columns
        return QtGui.QImage(cvMat.data, columns, rows, bytesPerLine, QtGui.QImage.Format_Indexed8)
    else:
        rows, columns, channels = cvMat.shape
        bytesPerLine = channels * columns
        return QtGui.QImage(cvMat.data, columns, rows, bytesPerLine, QtGui.QImage.Format_RGB888).rgbSwapped()


def base_to_cv2(b64str):
    try:
        datas = base64.b64decode(b64str.encode('utf-8'))
        datas = np.frombuffer(datas, np.uint8)
        image = cv2.imdecode(datas, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        if image is not None:
            return image
        else:
            raise ValueError("Failed to decode image")
    except Exception as e:
        raise ValueError("Error decoding image: " + str(e))


class ChooseConfigDialog(QDialog, Ui_ChooseConfigDialog):
    def __init__(self, parent=None):
        super(ChooseConfigDialog, self).__init__(parent)
        self.setupUi(self)
        # 配置列表
        self.config_list = []
        self.read_config_list()
        self.init_ui()
        self.init_event()

    def init_ui(self):
        self.config_listWidget.setCurrentRow(0)

    def init_event(self):
        """
        初始化事件
        :return:
        """
        # 选中内容改变事件
        self.config_listWidget.currentRowChanged.connect(self.config_list_widget_current_row_changed)

    def read_config_list(self):
        """
        读取配置列表
        :return:
        """
        self.config_listWidget.clear()
        for file in os.listdir('config'):
            if file.split('.')[-1] == 'json':
                self.config_listWidget.addItem(f'{file}')

    def config_list_widget_current_row_changed(self):
        """
        配置列表选中改变事件
        :return:
        """
        if self.config_listWidget.currentItem() is None:
            return
        file_name = self.config_listWidget.currentItem().text()
        with open(f'config/{file_name}', 'r') as f:
            content = f.read()
            content = json.dumps(json.loads(content), indent=4, ensure_ascii=False)
        self.show_config_json_textBrowser.setText(content)


class MainFrameShow(Ui_MainWindow):
    def __init__(self, MainWindow):
        super(MainFrameShow, self).setupUi(MainWindow)
        self.MainWindow = MainWindow
        self.pic_format = ['bmp', 'jpg', 'png', 'tif', 'gif', 'jpeg']
        self.dir_path = ''
        self.current_pic_path = ''
        self.crop_points = set()
        self.config = {
            'cut_width': 50,
            'cut_height': 50,
            'save_path': '',
            'read_path': '',
            'detect_model': 'ddddocr',
            'filter_perimeter_max': '',
            'filter_perimeter_min': '',
            'width_height_proportion_min': '',
            'width_height_proportion_max': '',
            'save_cut_img_check_box_state': True,
            'save_yolo_txt_check_box_state': True,
            'single_label_checkBox_state': True,
            'not_detect_checkBox': False
        }
        self.save_path = ''
        self.read_path = ''
        self.scene = QGraphicsScene()
        self.graphicsView.setScene(self.scene)
        # 设置scene为无法改变大小
        # self.graphicsView.setFixedSize(570, 511)  # 固定宽度为 800，高度为 600
        self.graphicsView.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)

        # 设置为完全填充
        self.graphicsView.setRenderHint(QtGui.QPainter.Antialiasing, True)
        # 默认保存目录按钮禁止
        self.save_config_btn.setEnabled(False)
        self.setting_save_dir_btn.setEnabled(False)
        self.batch_detect_btn.setEnabled(False)
        # onnx 模型推理
        self.detect_model_onnx_det = None
        self.init_ui()

    def init_ui(self):
        """
        初始化ui
        :return:
        """
        self.pushButton.clicked.connect(self.on_push_button_clicked)
        # self.fileListWidget.itemDoubleClicked.connect(self.file_list_widget_double_clicked)
        self.fileListWidget.currentRowChanged.connect(self.file_list_widget_double_clicked)

        # 在初始化时保存原始的 mousePressEvent
        self.original_mousePressEvent = self.graphicsView.mousePressEvent
        self.graphicsView.mousePressEvent = self.img_show_clicked

        self.save_config_btn.clicked.connect(self.save_config)
        self.setting_save_dir_btn.clicked.connect(self.setting_save_dir_btn_clicked)
        self.detect_model_comboBox.currentIndexChanged.connect(self.detect_model_combo_box_changed)
        # 连接右键菜单
        self.fileListWidget.setContextMenuPolicy(Qt.CustomContextMenu)
        self.fileListWidget.customContextMenuRequested.connect(self.file_list_widget_right_menu)
        # 设置自动推理按钮
        self.batch_detect_btn.clicked.connect(self.auto_predict_btn_clicked)
        # 设置关闭事件
        self.MainWindow.closeEvent = self.close_event
        # 设置快速读入配置按钮
        self.read_config_list_btn.clicked.connect(self.read_config_list_btn_clicked)
        # 设置按配置分割
        self.cut_by_config_btn.clicked.connect(self.cut_by_config_btn_clicked)
        # 设置简单分割
        self.cut_by_simple_btn.clicked.connect(self.cut_by_simple_btn_clicked)
        # 设置热键s
        # self.hot_key_s = QAction()
        sq = QShortcut(QKeySequence('s'), self.MainWindow)  # 通过QShortcut控件设置快捷方式，触发槽函数
        sq.activated.connect(self.next_img)
        # 设置热键w
        wq = QShortcut(QKeySequence('w'), self.MainWindow)  # 通过QShortcut控件设置快捷方式，触发槽函数
        wq.activated.connect(self.prev_img)

        # 扫描models目录下的所有模型文件
        for file in os.listdir('models'):
            if file.split('.')[-1] == 'onnx':
                self.detect_model_comboBox.addItem(file)
        qw = QWidget()
        DataCollectWidgetFunctionStarter(qw)
        self.tabWidget.addTab(qw, "数采区")

    def next_img(self):
        """
        下一张图片
        :return:
        """
        if self.fileListWidget.currentRow() == self.fileListWidget.count() - 1:
            return
        self.fileListWidget.setCurrentRow(self.fileListWidget.currentRow() + 1)

    def prev_img(self):
        """
        上一张图片
        :return:
        """
        if self.fileListWidget.currentRow() == 0:
            return
        self.fileListWidget.setCurrentRow(self.fileListWidget.currentRow() - 1)

    def read_config_list_btn_clicked(self):
        """
        读取配置列表按钮点击事件，打开一个dialog
        :return:
        """
        dialog = ChooseConfigDialog()
        if dialog.exec_() == QDialog.Accepted:
            # 用户点击了 OK 按钮
            logger.info("对话框接受")
            # 取选中的配置文件
            file_name = dialog.config_listWidget.currentItem().text()
            logger.info(f'选中的配置文件：{file_name}')
            with open(f'config/{file_name}', 'r') as f:
                self.config = json.load(f)
            if self.config.get('read_path') is None:
                QMessageBox.warning(self.MainWindow, '警告', '配置文件中没有read_path配置')
                return
            self.change_dir(self.config['read_path'])
        else:
            # 用户点击了 Cancel 按钮
            logger.info("取消选择")

    def auto_predict_btn_clicked(self):
        """
        自动推理按钮点击事件
        :return:
        """
        # 判断 save_path 是否为空，为空则提示
        if self.save_path == '':
            QMessageBox.warning(self.MainWindow, '警告', '请先设置保存目录')
            return
        # 判断文件量
        if self.fileListWidget.count() == 0:
            QMessageBox.warning(self.MainWindow, '警告', '当前目录下没有图片')
            return
        # 设置不选中任何文件
        self.fileListWidget.setCurrentRow(-1)
        # 遍历所有文件
        for i in range(self.fileListWidget.count()):
            self.next_img()
            QApplication.processEvents()

    def init_default_config(self):
        """
        恢复默认配置
        :return:
        """
        self.config = {
            'cut_width': 50,
            'cut_height': 50,
            'save_path': '',
            'read_path': '',
            'detect_model': 'ddddocr',
            'filter_perimeter_max': '',
            'filter_perimeter_min': '',
            'width_height_proportion_min': '',
            'width_height_proportion_max': '',
            'save_cut_img_check_box_state': True,
            'save_yolo_txt_check_box_state': True,
            'single_label_checkBox_state': True,
            'not_detect_checkBox': False
        }
        self.save_path = ''
        self.read_path = ''
        # 清空场景
        self.scene.clear()
        self.current_pic_path = ''
        self.crop_points = set()
        self.preview_cut_img_label.clear()

    def detect_model_combo_box_changed(self):
        """
        模型选择下拉框变化事件
        :return:
        """
        current_model = self.detect_model_comboBox.currentText()
        if current_model == 'ddddocr':
            self.detect_model_onnx_det = None
        else:
            self.detect_model_onnx_det = YOLOSystem(f'models/{current_model}', [str(i) for i in range(10000)])

    def init_and_read_config(self):
        """
        读取配置文件
        :return:
        """
        # 存储当前目录的配置，文件名称为当前目录信息做md5
        md5 = hashlib.md5()
        md5.update(self.dir_path.encode('utf-8'))
        url_md5 = md5.hexdigest()
        config_name = f'config/{url_md5}.json'
        if not os.path.exists(config_name):
            with open(config_name, 'w') as f:
                f.write(json.dumps(self.config, indent=4))
        else:
            with open(config_name, 'r') as f:
                self.config = json.load(f)
        #  设置cut_width
        self.cut_width.setText(str(self.config['cut_width']))
        # 设置cut_height
        self.cut_height.setText(str(self.config['cut_height']))
        # 设置save_path
        self.save_path = self.config['save_path']
        # 设置读入路径
        # self.read_path = self.config['read_path']
        # 设置过滤周长最大值
        self.filter_perimeter_max_edit.setText(self.config.get('filter_perimeter_max', ''))
        # 设置过滤周长最小值
        self.filter_perimeter_min_edit.setText(self.config.get('filter_perimeter_min', ''))
        # 设置宽高比例最小值
        self.width_height_proportion_min_edit.setText(self.config.get('width_height_proportion_min', ''))
        # 设置宽高比例最大值
        self.width_height_proportion_max_edit.setText(self.config.get('width_height_proportion_max', ''))
        # 设置模型选择下拉框
        index = self.detect_model_comboBox.findText(self.config.get('detect_model', 'ddddocr'))
        self.detect_model_comboBox.setCurrentIndex(index)
        # 设置是否保存切图
        self.save_cut_img_check_box.setChecked(self.config.get('save_cut_img_check_box_state', True))
        # 设置是否保存yolo txt文件
        self.save_yolo_txt_check_box.setChecked(self.config.get('save_yolo_txt_check_box_state', True))
        # 设置是否单标签模式
        self.single_label_checkBox.setChecked(self.config.get('single_label_checkBox_state', True))
        # 设置是否不检测
        self.not_detect_checkBox.setChecked(self.config.get('not_detect_checkBox', False))
        # 判断目录是否存在，不存在则创建
        if self.save_path != '' and not os.path.exists(self.save_path):
            os.makedirs(self.save_path)

    def save_config(self):
        """
        保存配置文件
        :return:
        """
        md5 = hashlib.md5()
        md5.update(self.dir_path.encode('utf-8'))
        url_md5 = md5.hexdigest()
        config_name = f'config/{url_md5}.json'
        self.config['cut_width'] = int(self.cut_width.text())
        self.config['cut_height'] = int(self.cut_height.text())
        self.config['save_path'] = self.save_path
        self.config['read_path'] = self.read_path
        self.config['detect_model'] = self.detect_model_comboBox.currentText()
        self.config['filter_perimeter_max'] = self.filter_perimeter_max_edit.text()
        self.config['filter_perimeter_min'] = self.filter_perimeter_min_edit.text()
        self.config['width_height_proportion_min'] = self.width_height_proportion_min_edit.text()
        self.config['width_height_proportion_max'] = self.width_height_proportion_max_edit.text()
        self.config['save_cut_img_check_box_state'] = self.save_cut_img_check_box.isChecked()
        self.config['save_yolo_txt_check_box_state'] = self.save_yolo_txt_check_box.isChecked()
        self.config['single_label_checkBox_state'] = self.single_label_checkBox.isChecked()
        self.config['not_detect_checkBox'] = self.not_detect_checkBox.isChecked()
        with open(config_name, 'w') as f:
            f.write(json.dumps(self.config, indent=4))
        QMessageBox.information(self.MainWindow, '提示', '保存成功')

    def check_config_is_change(self):
        """
        检查配置是否有改变
        :return:
        """
        # 读取配置文件
        md5 = hashlib.md5()
        md5.update(self.dir_path.encode('utf-8'))
        url_md5 = md5.hexdigest()
        config_name = f'config/{url_md5}.json'
        if not os.path.exists(config_name):
            return False
        with open(config_name, 'r') as f:
            config = json.load(f)
        if config['cut_width'] != int(self.cut_width.text()) or config['cut_height'] != int(self.cut_height.text()):
            return True

    def on_push_button_clicked(self):
        """
        选择目录按钮点击事件
        :return:
        """
        # 打开文件对话框，选择一个目录
        dir_path = QFileDialog.getExistingDirectory(self.MainWindow, "选择图片读取目录", "./")
        if not dir_path:
            return
        self.change_dir(dir_path)

    def reset_file_list_widget(self):
        """
        清空列表,并重新读取
        :return:
        """
        #
        self.fileListWidget.clear()
        # 遍历当前目录下的所有图片文件
        for file in sorted(os.listdir(self.dir_path)):
            if file.split('.')[-1] in self.pic_format:
                logger.info(f'{self.dir_path}/{file}')
                self.fileListWidget.addItem(f'{file}')

    def change_dir(self, dir_path: str):
        """
        更改目录
        :param dir_path:  更改的目录路径
        :return:
        """
        self.read_path = dir_path
        if self.dir_path == dir_path:
            self.reset_file_list_widget()
            return
        if self.dir_path != dir_path and self.dir_path != '' and self.check_config_is_change():
            # 说明更新了目录，需要询问是否保存配置
            reply = QMessageBox.question(self.MainWindow, '提示', '检测到您更换了目录，是否保存当前目录配置?',
                                         QMessageBox.Yes | QMessageBox.No,
                                         QMessageBox.No)
            if reply == QMessageBox.Yes:
                self.save_config()
        # 初始化默认配置，同时清空画布
        self.init_default_config()
        self.dir_path = dir_path
        self.read_path = dir_path
        self.init_and_read_config()
        logger.info(self.dir_path)
        # 设置 statusbar 文本
        self.statusbar.showMessage(f'数据目录：{self.dir_path}, 切图存储目录：{self.save_path}')
        # 清空列表
        self.reset_file_list_widget()
        # 保存配置按钮设置为可用
        self.save_config_btn.setEnabled(True)
        self.setting_save_dir_btn.setEnabled(True)
        if self.save_path != '':
            # 如果保存目录不为空，则设置为可用
            self.batch_detect_btn.setEnabled(True)

    def setting_save_dir_btn_clicked(self):
        """
        设置保存目录按钮点击事件
        :return:
        """
        # 打开文件对话框，选择一个目录
        dir_path = QFileDialog.getExistingDirectory(self.MainWindow, "选择图片保存目录", "./")
        if not dir_path:
            return
        self.save_path = dir_path
        self.statusbar.showMessage(f'数据目录：{self.dir_path}, 切图存储目录：{self.save_path}')
        QMessageBox.information(self.MainWindow, '提示', '设图片保存目录置成功')
        if self.save_path != '':
            # 如果保存目录不为空，则设置为可用
            self.batch_detect_btn.setEnabled(True)

    def file_list_widget_double_clicked(self):
        """
        文件列表双击事件
        :return:
        """
        # 判断 save_path 是否为空，为空则提示
        if self.save_path == '':
            QMessageBox.warning(self.MainWindow, '警告', '请先设置保存目录')
            return
        if self.fileListWidget.currentItem() is None:
            return
        # 获取当前选中的文件名
        file_name = self.fileListWidget.currentItem().text()
        logger.info(file_name)
        # 清空坐标列表
        self.crop_points = set()
        # 预览图清空
        self.preview_cut_img_label.clear()
        # 显示图片
        self.current_pic_path = f'{self.dir_path}/{file_name}'
        # 判断是否需要重新检测
        conf_dir = self.check_and_mkdir_conf_dir()
        current_img_conf = f'{conf_dir}/{self.current_pic_path.split("/")[-1].split(".")[0]}.conf'
        re_detect_flag = not os.path.exists(current_img_conf)
        self.show_pic_flow(re_detect_flag)

    def detect_pic(self, re_detect=False):
        """
        推理图片
        :param re_detect: 是否重新检测
        :return:
        """
        original_image = cv2.imread(self.current_pic_path)
        # 临时存储推理坐标
        temp_poses = []
        # 判断是否要重新推理
        if re_detect:
            if self.detect_model_onnx_det is not None:
                # onnx模型推理
                base64_content = base64.b64encode(open(self.current_pic_path, 'rb').read()).decode()
                onnx_det_result = self.detect_model_onnx_det.serving_method(base64_content)
                logger.info(f'onnx_det_result: {onnx_det_result}')
                for res in onnx_det_result:
                    if res.get('crop') is not None:
                        crop = res.get('crop')
                        crop.append(res.get('classes', '0'))
                        logger.info(f'crop: {crop}')
                        temp_poses.append(crop)
            else:
                # ddddocr根据原图识别
                dddd = ddddocr.DdddOcr(det=True, show_ad=False)
                temp_poses = dddd.detection(open(self.current_pic_path, 'rb').read())
        else:
            # 不需要重新推理，则把之前的坐标赋值给temp_poses
            for crop in self.crop_points:
                temp_poses.append(list(crop))
        for index, pose in enumerate(temp_poses):
            # 检测区域，如果区域为负数，则为0
            if len(pose) == 4:
                pose.append('0')
            x1, y1, x2, y2, classes = pose
            x1 = max(0, x1)
            y1 = max(0, y1)
            # 如果区域超出图片范围，则设置为图片的最大值
            x2 = min(max(0, x2), original_image.shape[1])
            y2 = min(max(0, y2), original_image.shape[0])
            # 计算周长
            perimeter = (x2 - x1) * 2 + (y2 - y1) * 2
            if (self.filter_perimeter_max_edit.text() != ''
                and perimeter > int(self.filter_perimeter_max_edit.text())) \
                    or (self.filter_perimeter_min_edit.text() != ''
                        and perimeter < int(self.filter_perimeter_min_edit.text())):
                continue
            # 计算宽高比例
            width_height_proportion = (x2 - x1) / (y2 - y1)
            if (self.width_height_proportion_max_edit.text() != '' and width_height_proportion > float(
                    self.width_height_proportion_max_edit.text())) or (
                    self.width_height_proportion_min_edit.text() != '' and width_height_proportion < float(
                self.width_height_proportion_min_edit.text())):
                continue
            # 坐标存储至self.crop_points，用于后续
            self.crop_points.add((x1, y1, x2, y2, classes))
            # 判断是否需要保存图片,如果需要保存，则操作一下
            if self.save_cut_img_check_box.isChecked():
                # 保存目录
                save_dir = self.save_path + '/' + self.current_pic_path.split('/')[-1].split('.')[0]
                logger.info(f'save_dir: {save_dir}')
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                # 删除x1_y1_x2_y2.jpg或者x1_y1_x2_y2_0.jpg的图片
                if os.path.exists(f'{save_dir}/{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}.jpg'):
                    os.remove(f'{save_dir}/{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}.jpg')
                if os.path.exists(f'{save_dir}/{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}_0.jpg'):
                    os.remove(f'{save_dir}/{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}_0.jpg')
                # 保存图片
                cv2.imwrite(f'{save_dir}/{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}_{classes}.jpg',
                            original_image[int(y1):int(y2), int(x1):int(x2)])

    def show_pic_flow(self, re_detect=False):
        """
        拆出展示图片的流程，所有图片都需要经过先推理，再展示
        推理主要是为了获取坐标，保存到全局
        展示主要是为了展示图片
        画框逻辑是为了展示图片时识别框
        :return:
        """
        original_image = cv2.imread(self.current_pic_path)
        # 图片和 label 的宽高比例
        w = self.graphicsView.width()
        h = self.graphicsView.height()
        # 图片缩放比例
        w_scale = w / original_image.shape[1]
        h_scale = h / original_image.shape[0]
        # 0.初始化配置
        self.init_conf()
        # 1. 推理图片，拿到坐标
        self.detect_pic(re_detect)
        # 2. 展示图片
        self.show_pic()
        # 3. 重新画框
        self.re_draw(w_scale, h_scale)
        # 4. 保存yolo信息
        self.save_yolo_model_txt()
        # 保存框框信息
        self.save_pic_crop_config(w_scale, h_scale, self.current_pic_path)
        # 随便选一个区域，预览小图
        if len(self.crop_points) > 0:
            x1, y1, x2, y2, classes = list(self.crop_points)[0]
            self.show_small_pic(x1, y1, x2, y2)

    def init_conf(self):
        """
        初始化一张图片的区域配置信息，会读入到crop_points中
        :return:
        """
        # 判断当前图片是否有conf目录
        conf_dir = self.check_and_mkdir_conf_dir()
        current_img_conf = f'{conf_dir}/{self.current_pic_path.split("/")[-1].split(".")[0]}.conf'
        if not os.path.exists(current_img_conf):
            # 如果不存在，说明是第一次处理，需要进行模型推理
            if self.not_detect_checkBox.isChecked():
                # 如果不检测，直接显示
                self.crop_points = set()
            else:
                self.detect_pic()
        else:
            # 如果存在，说明已经处理过，直接显示，
            with open(current_img_conf, 'r') as f:
                poses = f.readlines()
                for pose in poses:
                    pose = pose.strip()
                    if pose == '':
                        continue
                    x1, y1, x2, y2, classes = pose.split(' ')
                    self.crop_points.add((float(x1), float(y1), float(x2), float(y2), classes))

    def check_and_mkdir_conf_dir(self):
        """
        检查并创建当前图片的conf目录
        :return:  conf目录地址
        """
        read_path = self.read_path
        if read_path[-1] == '/':
            read_path = read_path[:-1]
        last_dirname = read_path.split('/')[-1]
        conf_dir = f'{self.read_path}/{last_dirname}_conf'
        if not os.path.exists(conf_dir):
            os.makedirs(conf_dir)
        return conf_dir

    def show_pic(self):
        """
        显示图片
        :return:
        """
        if self.current_pic_path == '':
            QMessageBox.warning(self.MainWindow, '警告', '请先选择一张图片')
        # 图片展示到label
        image = cv2.imread(self.current_pic_path)
        image = cv2.resize(image, (self.graphicsView.width(), self.graphicsView.height()))
        # 重新绘制图片
        self.graphicsView_re_draw(image)

    def show_small_pic(self, x1, y1, x2, y2):
        """
        预览小图
        :param x1:
        :param y1:
        :param x2:
        :param y2:
        :return:
        """
        im = cv2.imread(self.current_pic_path)
        im = im[int(y1):int(y2), int(x1):int(x2)]
        im = cv2.resize(im, (self.preview_cut_img_label.width(), self.preview_cut_img_label.height()))
        image = CvMatToQImage(im)
        pixmap = QPixmap.fromImage(image)
        self.preview_cut_img_label.setPixmap(pixmap)

    def save_yolo_model_txt(self):
        """
        保存yolo模型的txt文件
        :return:
        """
        content_list = []
        # 取原图的大小
        image = cv2.imread(self.current_pic_path)
        w, h = image.shape[1], image.shape[0]
        model = '{} {} {} {} {}'
        for pos in self.crop_points:
            # 根据原图的坐标，换算成比例
            cx = (pos[0] + pos[2]) / 2 / w
            cy = (pos[1] + pos[3]) / 2 / h
            w_ = (pos[2] - pos[0] + 8) / w
            h_ = (pos[3] - pos[1] + 8) / h
            content_list.append(model.format(pos[4], cx, cy, w_, h_))
        with open(f'{self.dir_path}/{self.current_pic_path.split("/")[-1].split(".")[0]}.txt', 'w') as f:
            f.write('\n'.join(content_list))

    def img_show_clicked(self, event: QMouseEvent):
        """
        图片点击事件
        会在图片上画一个框，表示选中的区域，框的大小和设置的大小有关
        :return:
        """
        if not self.current_pic_path:
            return
        # 取点击的坐标
        click_x = event.x()
        click_y = event.y()
        logger.info(f'click_x: {click_x}, click_y: {click_y}')
        # 获取当前图片展示区的大小
        w = self.graphicsView.width()
        h = self.graphicsView.height()
        # 获取当前图片
        image = cv2.imread(self.current_pic_path)
        # 备份原图，用于计算比例
        original_image = image.copy()
        # 缩放图片
        image = cv2.resize(image, (w, h))
        # 图片缩放比例
        w_scale = w / original_image.shape[1]
        h_scale = h / original_image.shape[0]
        item = self.graphicsView.itemAt(event.pos())  # 获取点击的项
        # 判断是否是ResizableRect
        if isinstance(item, ResizableRect) and event.button() == Qt.LeftButton:
            # 如果是子项，传递事件给子项处理，不干扰
            logger.info("Mouse press on ResizableRect")
            # 将事件传递给 ResizableRect 自己处理
            self.original_mousePressEvent(event)
            return  # 直接返回，不执行后续自定义逻辑
        else:
            if event.button() == Qt.LeftButton:
                # 设置框的大小，这是原图的大小
                box_w = int(self.cut_width.text())
                box_h = int(self.cut_height.text())
                # 换算成缩放后的大小
                box_w = math.floor(box_w * w_scale)
                box_h = math.floor(box_h * h_scale)
                # 获取框的位置
                x1 = click_x - box_w // 2
                y1 = click_y - box_h // 2
                x2 = click_x + box_w // 2
                y2 = click_y + box_h // 2
                classes = '0'
                if not self.single_label_checkBox.isChecked():
                    # 弹框输入标签
                    classes, ok = QInputDialog.getText(self.MainWindow, '标签', '请输入标签')
                    if not ok:
                        return
                if x1 < 0:
                    x1 = 0
                    x2 = box_w
                if x2 > w:
                    x1 = w - box_w
                    x2 = w
                if y1 < 0:
                    y1 = 0
                    y2 = box_h
                if y2 > h:
                    y1 = h - box_h
                    y2 = h
                # 换算成原图比例
                original_x1 = math.floor(x1 * original_image.shape[1] / w)
                original_x2 = math.floor(x2 * original_image.shape[1] / w)
                original_y1 = math.floor(y1 * original_image.shape[0] / h)
                original_y2 = math.floor(y2 * original_image.shape[0] / h)
                # 预览小图
                self.show_small_pic(original_x1, original_y1, original_x2, original_y2)
                self.crop_points.add((original_x1, original_y1, original_x2, original_y2, classes))
                if self.save_cut_img_check_box.isChecked():
                    # 如果选中了，那么就保存推理图片
                    cv2.imwrite(
                        f'{self.save_path}/{self.current_pic_path.split("/")[-1].split(".")[0]}/{original_x1}_{original_y1}_{original_x2}_{original_y2}_{classes}.jpg',
                        original_image[original_y1:original_y2, original_x1:original_x2])
                self.graphicsView_re_draw(image)
                # 画上所有的框
                self.re_draw(w_scale, h_scale)
                # 重置 yolo 模型的txt文件
                self.save_yolo_model_txt()
                # 保存框框信息
                self.save_pic_crop_config(w_scale, h_scale, self.current_pic_path)
            elif event.button() == Qt.RightButton:
                # 根据缩放比例，还原成原图的坐标
                original_x = math.floor(click_x / w_scale)
                original_y = math.floor(click_y / h_scale)
                # 删除对应区域
                self.crop_points = {p for p in self.crop_points if
                                    not (p[0] < original_x < p[2] and p[1] < original_y < p[3])}
                if not isinstance(item, ResizableRect):
                    # 如果不是 ResizableRect，说明是在空白处点击，不管他
                    return
                # 重新切图
                self.re_write_cut_icon()
                # 清理掉所有的框
                self.just_clean_scene_resizable_rect()
                # 重新画框
                self.re_draw(w_scale, h_scale)
                # 重置 yolo 模型的txt文件
                self.save_yolo_model_txt()
                # 保存框框信息
                self.save_pic_crop_config(w_scale, h_scale, self.current_pic_path)
            else:
                event.ignore()
        self.original_mousePressEvent(event)
        QApplication.processEvents()

    def re_write_cut_icon(self):
        """
        重新写入切图
        :return:
        """
        # 获取当前图片
        image = cv2.imread(self.current_pic_path)
        save_dir = f'{self.save_path}/{self.current_pic_path.split("/")[-1].split(".")[0]}'
        # 删除所有图片
        if os.path.exists(save_dir):
            for file in os.listdir(save_dir):
                # 删除file
                os.remove(f'{save_dir}/{file}')
        # 重新写入
        for crop in self.crop_points:
            x1, y1, x2, y2, classes = crop
            cv2.imwrite(f'{save_dir}/{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}_{classes}.jpg',
                        image[int(y1):int(y2), int(x1):int(x2)])

    def re_draw(self, w_scale, h_scale):
        """
        重新画框
        :param w_scale:  宽度缩放比例
        :param h_scale:  高度缩放比例
        :return:
        """
        # 计算均色
        border_color = self.calc_border_color()

        for index, pose in enumerate(self.crop_points):
            # 添加可调整大小的矩形
            rect_item = ResizableRect(QRectF(
                pose[0] * w_scale,
                pose[1] * h_scale,
                pose[2] * w_scale - pose[0] * w_scale,
                pose[3] * h_scale - pose[1] * h_scale),
                self.graphicsView.width(),
                self.graphicsView.height(),
                mouse_release_event_hook=self.save_pic_crop_config,
                mouse_move_event_hook=self.show_small_pic,
                pic_path=self.current_pic_path,
                w_scale=w_scale,
                h_scale=h_scale,
                no_hover_pen_color=border_color)
            self.scene.addItem(rect_item)
            # 设置在顶层
            rect_item.setZValue(2000 + index)
        QApplication.processEvents()

    def calc_border_color(self) -> QColor:
        """
        计算边框颜色
        通过背景图的均色来计算边框颜色，使其更加醒目
        取的是均色的反色
        :return:
        """
        items = self.scene.items()
        average_color = None
        for item in items:
            if isinstance(item, QGraphicsPixmapItem):
                image = item.pixmap().toImage()
                # 转换为 NumPy 数组
                ptr = image.bits()
                ptr.setsize(image.byteCount())
                data = np.array(ptr).reshape(image.height(), image.width(), 4)  # RGBA
                # 计算均色
                average_color = data.mean(axis=(0, 1))[:3]  # 取 RGB 平均值
                break
        if average_color is not None:
            # average_color = 255 - average_color
            # 判断average_color和黑色接近还是和白色接近

            # 计算最醒目的颜色
            border_color = self.design_contrasting_color(average_color)
            border_color.setAlpha(140)
            return border_color
        return QColor(0, 0, 255, 140)

    @staticmethod
    def design_contrasting_color(mean_color) -> QColor:
        # 计算亮度
        brightness = 0.2126 * mean_color[0] + 0.7152 * mean_color[1] + 0.0722 * mean_color[2]
        # 判断接近的颜色
        if brightness < 128:  # 亮度小于128则接近黑色
            # 返回白色
            return QColor(255, 255, 255)
        else:  # 亮度大于或等于128则接近白色
            # 返回黑色
            return QColor(0, 0, 0)

    def graphicsView_re_draw(self, image: cv2.typing.MatLike):
        """
        graphicsView画布重新画图像
        :param image:
        :return:
        """
        # 显示图片
        image = CvMatToQImage(image)
        pixmap = QPixmap.fromImage(image)
        # 设置图片
        self.scene.clear()
        pixmap_item = self.scene.addPixmap(pixmap)
        self.graphicsView.fitInView(pixmap_item, Qt.KeepAspectRatio)
        QApplication.processEvents()

    def just_clean_scene_resizable_rect(self):
        """
        清空scene中的ResizableRect
        :return:
        """
        for item in self.scene.items():
            if isinstance(item, ResizableRect):
                self.scene.removeItem(item)
        QApplication.processEvents()

    def file_list_widget_right_menu(self, pos):
        """
        文件列表右键菜单，主要用于重新推理
        :param pos:
        :return:
        """
        menu = QMenu(self.MainWindow)
        menu.addAction(QAction('重新推理', self.MainWindow))
        # 菜单事件处理
        action = menu.exec_(self.fileListWidget.mapToGlobal(pos))
        if action:
            self.current_pic_path = f'{self.dir_path}/{self.fileListWidget.currentItem().text()}'
            if action.text() == '重新推理':
                if not self.not_detect_checkBox.isChecked():
                    # self.detect_pic()
                    # self.show_pic()
                    self.show_pic_flow(re_detect=True)

    def save_pic_crop_config(self, w_scale, h_scale, pic_path):
        """
        保存某图的切图信息
        :return:
        """
        # 获取graphview中的所有item
        logger.info('save_pic_crop_config')
        # 取ResizabelRect的信息
        crop_points = self.get_graphview_resizable_rect_point(w_scale, h_scale)
        logger.info(crop_points)
        # 更新self.crop_points
        self.crop_points = set()
        for crop in crop_points:
            self.crop_points.add(tuple(crop))
        # 检查并创建当前图片的conf目录
        conf_dir = self.check_and_mkdir_conf_dir()
        pic_conf_path = f'{conf_dir}/{pic_path.split("/")[-1].split(".")[0]}.conf'
        logger.info(f'自动保存目录: {pic_conf_path}')
        # 保存当前图片的切图信息
        with open(f'{pic_conf_path}', 'w') as f:
            lines = []
            for crop in crop_points:
                lines.append(' '.join([str(i) for i in crop]))
            f.write('\n'.join(lines))
        # 重新写入切图
        self.re_write_cut_icon()

    def get_graphview_resizable_rect_point(self, w_scale, h_scale):
        """
        获取graphview中的所有ResizableRect的坐标
        :return:
        """
        # 获取graphview中的所有item
        items = self.scene.items()
        logger.info('save_pic_crop_config')
        # 取ResizabelRect的信息
        crop_points = []
        for item in items:
            if isinstance(item, ResizableRect):
                rect = item.rect()
                # 取左上角的坐标
                top_left = item.mapToScene(item.rect().topLeft())  # 获取左上角在场景中的坐标
                x1 = top_left.x() / w_scale
                y1 = top_left.y() / h_scale
                x2 = x1 + rect.width() / w_scale
                y2 = y1 + rect.height() / h_scale
                crop_points.append([x1, y1, x2, y2, item.classes_name])
        return crop_points

    def cut_by_config_btn_clicked(self):
        """
        数据集根据配置切分按钮点击事件
        :return:
        """
        # 判断 save_path 是否为空，为空则提示
        if self.read_path == '':
            QMessageBox.warning(self.MainWindow, '警告', '请先设置读取目录')
            return
        # 判断文件量
        if len(os.listdir(self.read_path)) == 0:
            QMessageBox.warning(self.MainWindow, '警告', '图片目录为空')
            return
        # 确认是否要切分
        reply = QMessageBox.question(self.MainWindow, '提示',
                                     '按配置切分数据集:\n会根据上面的配置进行数据集拆分\n是否要按配置切分?如确定，则会清空输出目录后再生成!',
                                     QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.No:
            return
        DatasetsCutUtil(read_path=self.read_path,
                        train_rate=self.train_datasets_by_rate_spinBox.value() / 100,
                        train_num=self.train_datasets_by_nul_spinBox.value(),
                        cut_type=0,
                        cut_config_type=0 if self.rate_radioButton.isChecked() else 1)
        QMessageBox.information(self.MainWindow, '提示', '按配置切分完成')

    def cut_by_simple_btn_clicked(self):
        """
        数据集简单切分按钮点击事件
        :return:
        """
        # 判断 save_path 是否为空，为空则提示
        if self.read_path == '':
            QMessageBox.warning(self.MainWindow, '警告', '请先设置读取目录')
            return
        # 判断文件量
        if len(os.listdir(self.read_path)) == 0:
            QMessageBox.warning(self.MainWindow, '警告', '图片目录为空')
            return
        # 确认是否要切分
        reply = QMessageBox.question(self.MainWindow, '提示',
                                     '简单切分数据集:\n小于300张图片,训练集和测试集共用\n小于500张，按9:1切分\n大于500张，按8:2切分\n是否要简单切分?如确定，则会清空输出目录后再生成!',
                                     QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.No:
            return
        DatasetsCutUtil(read_path=self.read_path, cut_type=1)
        QMessageBox.information(self.MainWindow, '提示', '简单切分完成')

    # 设置销毁前的操作
    def close_event(self, event):
        if self.dir_path != '' and self.check_config_is_change():
            # 询问是否要保存配置
            reply = QMessageBox.question(self.MainWindow, '提示', '是否保存配置', QMessageBox.Yes | QMessageBox.No,
                                         QMessageBox.No)
            if reply == QMessageBox.Yes:
                self.save_config()
                event.accept()


if __name__ == '__main__':
    # 判断当前目录下是否存在 config目录
    if not os.path.exists('config'):
        os.makedirs('config')
    app = QApplication(sys.argv)
    MainWindow = QMainWindow()
    # 设置MainWindow不能调整大小
    MainWindow.setFixedSize(1070, 624)
    ui = MainFrameShow(MainWindow)
    MainWindow.show()
    sys.exit(app.exec_())
